from sqlalchemy import Column, String, Integer, Float
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from init.mysql import engine

Base = declarative_base()


class StockFund(Base):
    __tablename__ = 'stock_fund'
    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(String(6))
    jzrq = Column(String(10))
    jjmc = Column(String(255))
    jjdm = Column(String(20))
    ccsl = Column(Integer)
    ltbl = Column(Float)
    cgsz = Column(Float)
    jzbl = Column(Float)

    def __init__(self, code, jzrq, jjmc, jjdm, ccsl, ltbl, cgsz, jzbl):
        self.code = code
        self.jzrq = jzrq
        self.jjmc = jjmc
        self.jjdm = jjdm
        self.ccsl = ccsl
        self.ltbl = ltbl
        self.cgsz = cgsz
        self.jzbl = jzbl

    @staticmethod
    def insert_or_update(code, data_list):
        session = sessionmaker(bind=engine)()
        for data in data_list:
            existing_stock_fund = session.query(StockFund).filter_by(code=code, jzrq=data['jzrq']).first()
            if existing_stock_fund:
                existing_stock_fund.jjmc = data['jjmc']
                existing_stock_fund.jjdm = data['jjdm']
                existing_stock_fund.ccsl = data['ccsl']
                existing_stock_fund.ltbl = data['ltbl']
                existing_stock_fund.cgsz = data['cgsz']
                existing_stock_fund.jzbl = data['jzbl']
            else:
                session.add(
                    StockFund(code=code, jzrq=data['jzrq'], jjmc=data['jjmc'], jjdm=data['jjdm'], ccsl=data['ccsl'],
                              ltbl=data['ltbl'], cgsz=data['cgsz'], jzbl=data['jzbl']))
        session.commit()
        session.close()


Base.metadata.create_all(engine)
